package trace

import (
	"Demo/book/pkg/grpcx/Interceptors"
	"context"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"strconv"
)

type InterceptorBuilder struct {
	tracer     trace.Tracer
	propagator propagation.TextMapPropagator
	Interceptors.Builder
}

func (b *InterceptorBuilder) BuilderClient() grpc.UnaryClientInterceptor {
	propagation := b.propagator
	if propagation == nil {
		propagation = otel.GetTextMapPropagator()
	}
	tracer := b.tracer
	if tracer == nil {
		tracer = otel.Tracer("book/pkg/grpcx/interceptors/trace")
	}
	attrs := []attribute.KeyValue{
		semconv.RPCSystemKey.String("grpc"),
		attribute.Key("rpc.grpc.kind").String("unary"),
		attribute.Key("rpc.component").String("server"),
	}
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) (err error) {
		ctx, span := tracer.Start(ctx, method, trace.WithSpanKind(trace.SpanKindClient), trace.WithAttributes(attrs...))
		defer span.End()
		//inject 过程
		ctx = inject(ctx, propagation)
		defer func() {
			if err != nil {
				span.RecordError(err)
				st, _ := status.FromError(err)
				if st != nil {
					val, _ := strconv.Atoi(st.Code().String())
					span.SetAttributes(semconv.RPCGRPCStatusCodeKey.Int64(int64(val)))
				}
				span.SetStatus(codes.Error, err.Error())
			} else {
				span.SetStatus(codes.Ok, "OK")
			}
			span.End()
		}()
		err = invoker(ctx, method, req, reply, cc, opts...)
		return err
	}
}

func (b *InterceptorBuilder) BuilderServer() grpc.UnaryServerInterceptor {
	propagation := b.propagator
	if propagation == nil {
		propagation = otel.GetTextMapPropagator()
	}
	tracer := b.tracer
	if tracer == nil {
		tracer = otel.Tracer("book/pkg/grpcx/interceptors/trace")
	}
	attrs := []attribute.KeyValue{
		semconv.RPCSystemKey.String("grpc"),
		attribute.Key("rpc.grpc.kind").String("unary"),
		attribute.Key("rpc.component").String("server"),
	}
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		ctx = extract(ctx, propagation)
		ctx, span := tracer.Start(ctx, info.FullMethod, trace.WithSpanKind(trace.SpanKindServer), trace.WithAttributes(attrs...))

		defer span.End()

		span.SetAttributes(semconv.RPCMethodKey.String(info.FullMethod), semconv.NetPeerNameKey.String(b.PeerName(ctx)),
			attribute.Key("net.peer.ip").String(b.PeerIP(ctx)))
		resp, err := handler(ctx, req)
		defer func() {
			//就要结束了
			if err != nil {
				span.RecordError(err)
			} else {
				span.SetStatus(codes.Ok, "OK")
			}
		}()

		return resp, err
	}
}

func inject(ctx context.Context, propagators propagation.TextMapPropagator) context.Context {
	md, ok := metadata.FromOutgoingContext(ctx)
	if !ok {
		md = metadata.MD{}
	}
	//把元数据放回去ctx，具体怎么放，放什么内容，由propagators来决定，由zipkin和skywalking来决定
	propagators.Inject(ctx, GRPCHeaderCarrier(md))
	return metadata.NewOutgoingContext(ctx, md)
}

func extract(ctx context.Context, p propagation.TextMapPropagator) context.Context {
	//拿到客户端过来的链路元数据
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		md = metadata.New(map[string]string{})
	}
	//把这个me注入到 ctx 里面。
	return p.Extract(ctx, GRPCHeaderCarrier(md))
}

type GRPCHeaderCarrier metadata.MD

func (G GRPCHeaderCarrier) Get(key string) string {
	vals := metadata.MD(G).Get(key)
	if len(vals) > 0 {
		return vals[0]
	}
	return ""
}

func (G GRPCHeaderCarrier) Set(key string, value string) {
	metadata.MD(G).Set(key, value)
}

func (G GRPCHeaderCarrier) Keys() []string {
	keys := make([]string, 0, len(G))
	for _, key := range keys {
		keys = append(keys, key)
	}
	return keys
}
